#include <windows.h>
#include <shlobj.h>
#include <stdio.h>
#include <string.h>

#define MENU_KEY L"Directory\\Background\\shell\\PasteAsHardLink"
#define APP_NAME L"Paste as Hard Link"

// 注册右键菜单
BOOL RegisterContextMenu() {
    HKEY hKey;
    WCHAR modulePath[MAX_PATH];
    WCHAR command[MAX_PATH * 2]; // 确保有足够空间

    GetModuleFileNameW(NULL, modulePath, MAX_PATH);
    _snwprintf_s(command, _countof(command), _TRUNCATE, L"\"%s\" \"%%V\"", modulePath);

    if (RegCreateKeyExW(HKEY_CLASSES_ROOT, MENU_KEY, 0, NULL, 0, KEY_WRITE, NULL, &hKey, NULL) != ERROR_SUCCESS)
    {
        return FALSE;
    }

    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE*)APP_NAME, (wcslen(APP_NAME) + 1) * sizeof(WCHAR));

    HKEY hCommandKey;
    if (RegCreateKeyExW(hKey, L"command", 0, NULL, 0, KEY_WRITE, NULL, &hCommandKey, NULL) != ERROR_SUCCESS)
    {
        RegCloseKey(hKey);
        return FALSE;
    }

    RegSetValueExW(hCommandKey, NULL, 0, REG_SZ, (BYTE*)command, (wcslen(command) + 1) * sizeof(WCHAR));

    RegCloseKey(hCommandKey);
    RegCloseKey(hKey);
    return TRUE;
}

// 取消注册
BOOL UnregisterContextMenu() {
    return RegDeleteTreeW(HKEY_CLASSES_ROOT, MENU_KEY) == ERROR_SUCCESS;
}

// 创建硬链接
BOOL CreateHardLinkForFile(LPCWSTR source, LPCWSTR destFolder) {
    WCHAR destPath[MAX_PATH * 2]; // 增加缓冲区大小
    WCHAR fileName[MAX_PATH];

    // 从源路径提取文件名
    const WCHAR* lastBackslash = wcsrchr(source, L'\\');
    if (!lastBackslash) return FALSE;
    wcscpy_s(fileName, _countof(fileName), lastBackslash + 1);

    // 构建目标路径
    _snwprintf_s(destPath, _countof(destPath), _TRUNCATE, L"%s\\%s", destFolder, fileName);

    // 检查源文件是否存在
    if (GetFileAttributesW(source) == INVALID_FILE_ATTRIBUTES)
    {
        return FALSE;
    }

    // 创建硬链接
    return CreateHardLinkW(destPath, source, NULL);
}

// 处理粘贴操作
void ProcessPasteOperation(LPCWSTR targetDir) {
    if (!OpenClipboard(NULL))
    {
        MessageBoxW(NULL, L"无法访问剪贴板", L"错误", MB_ICONERROR);
        return;
    }

    HDROP hDrop = (HDROP)GetClipboardData(CF_HDROP);
    if (!hDrop)
    {
        CloseClipboard();
        MessageBoxW(NULL, L"剪贴板中没有文件", L"错误", MB_ICONINFORMATION);
        return;
    }

    UINT fileCount = DragQueryFileW(hDrop, 0xFFFFFFFF, NULL, 0);
    BOOL success = TRUE;

    for (UINT i = 0; i < fileCount; i++)
    {
        WCHAR sourceFile[MAX_PATH * 2]; // 增加缓冲区大小
        DragQueryFileW(hDrop, i, sourceFile, _countof(sourceFile));

        if (!CreateHardLinkForFile(sourceFile, targetDir))
        {
            WCHAR errorMsg[512];
            _snwprintf_s(errorMsg, _countof(errorMsg), _TRUNCATE, 
                L"创建硬链接失败: %s\n错误代码: %d", sourceFile, GetLastError());
            MessageBoxW(NULL, errorMsg, L"错误", MB_ICONERROR);
            success = FALSE;
            break;
        }
    }

    if (success && fileCount > 0)
    {
        MessageBoxW(NULL, L"硬链接创建成功", L"完成", MB_ICONINFORMATION);
    }

    DragFinish(hDrop);
    CloseClipboard();
}

int wmain(int argc, WCHAR* argv[]) {
    if (argc == 2)
    {
        // 检查是否为注册/注销命令
        if (wcscmp(argv[1], L"/register") == 0)
        {
            if (RegisterContextMenu())
            {
                MessageBoxW(NULL, L"菜单注册成功", L"成功", MB_ICONINFORMATION);
            }
            else
            {
                MessageBoxW(NULL, L"菜单注册失败", L"错误", MB_ICONERROR);
            }
        }
        else if (wcscmp(argv[1], L"/unregister") == 0)
        {
            if (UnregisterContextMenu())
            {
                MessageBoxW(NULL, L"菜单已移除", L"成功", MB_ICONINFORMATION);
            }
            else
            {
                MessageBoxW(NULL, L"菜单移除失败", L"错误", MB_ICONERROR);
            }
        }
        else
        {
            // 正常调用：传递了目标目录
            ProcessPasteOperation(argv[1]);
        }
    }
    else
    {
        WCHAR msg[512];
        _snwprintf_s(msg, _countof(msg), _TRUNCATE,
            L"用法:\n"
            L"  %s [目标文件夹]   - 粘贴文件 \n"
            L"  %s /register    - 注册菜单\n"
            L"  %s /unregister  - 注销菜单",
            argv[0], argv[0], argv[0]);
        MessageBoxW(NULL, msg, L"帮助", MB_ICONINFORMATION);
    }

    return 0;
}

// 提供入口点以支持同时运行ANSI和Unicode环境
int main(int argc, char* argv[]) {
    LPWSTR *wargv;
    int wargc;
    wargv = CommandLineToArgvW(GetCommandLineW(), &wargc);
    int result = wmain(wargc, wargv);
    LocalFree(wargv);
    return result;
}
